from cozepy import COZE_CN_BASE_URL, Coze, TokenAuth, Stream, WorkflowEvent, WorkflowEventType


def run_coze_workflow(user_input: str, workflow_id: str, api_token: str, base_url: str = COZE_CN_BASE_URL):
    """
    运行Coze工作流并返回消息事件内容

    参数:
        user_input: 用户输入内容
        workflow_id: 工作流ID
        api_token: Coze API访问令牌
        base_url: Coze API基础URL(默认为中国区)

    返回:
        List[str]: 收到的所有消息内容
    """
    # 初始化Coze客户端
    coze = Coze(auth=TokenAuth(token=api_token), base_url=base_url)

    # 用于收集结果
    messages = []

    def handle_workflow_iterator(stream: Stream[WorkflowEvent]):
        """处理工作流事件迭代器"""
        for event in stream:
            if event.event == WorkflowEventType.MESSAGE:
                messages.append(event.message.content)
            elif event.event == WorkflowEventType.ERROR:
                messages.append(f"[错误] {event.error.message}")
            elif event.event == WorkflowEventType.INTERRUPT:
                handle_workflow_iterator(
                    coze.workflows.runs.resume(
                        workflow_id=workflow_id,
                        event_id=event.interrupt.interrupt_data.event_id,
                        resume_data="hey",
                        interrupt_type=event.interrupt.interrupt_data.type,
                    )
                )

    # 运行工作流并收集结果
    handle_workflow_iterator(
        coze.workflows.runs.stream(
            workflow_id=workflow_id,
            parameters={"input": user_input}
        )
    )

    return messages


# 使用示例
if __name__ == "__main__":
    # 配置参数
    WORKFLOW_ID = "7523894481670881331"
    API_TOKEN = "pat_JzRhQ45cdVnjzwjPLHCJF1MmYI06vuCCFaaFBW65WkBCB4l9lvWhtJWOMxd3LOM1"

    # 运行工作流
    results = run_coze_workflow("生成一张猫的图片", WORKFLOW_ID, API_TOKEN)
    print(results)
    for msg in results:
        print(msg)